
from copy import deepcopy
import math
import numpy as np
import os
import cv2
from PIL import Image
import random
from scipy.ndimage import zoom
import torch
from torch.utils.data import Dataset
from torchvision import transforms
from dataset.transform import random_flip, random_rotate, blur, obtain_cutmix_box


class CarotidArteryDataset(Dataset):
    def __init__(self, name, mode, cls_name, root, u_root=None, size=None, id_path=None, nsample=None):
        self.name = name
        self.mode = mode
        self.cls_name = cls_name
        self.root = root
        self.u_root = u_root
        self.size = size

        if mode == 'train_l' or mode == 'train_u':
            with open(id_path, 'r', encoding='utf-8') as f:
                self.ids = f.read().splitlines()
                print(f"Original {self.mode} samples: {len(self.ids)}")
            if mode == 'train_u':
                self.ids = self.ids
                print(f"Downsampled {self.mode} samples: {len(self.ids)}")
            if mode == 'train_l' and nsample is not None:
                self.ids *= math.ceil(nsample / len(self.ids))
                self.ids = self.ids[:nsample]
                print(f"Upsampled {self.mode} samples: {len(self.ids)}")
        else:
            with open('splits/%s/valid.txt' % name, 'r', encoding='utf-8') as f:
                self.ids = f.read().splitlines()
                print(f"{self.mode} samples: {len(self.ids)}")

    def __getitem__(self, item):
        id = self.ids[item]
        if self.mode == 'train_l' or self.mode == 'val':
            img = Image.open(os.path.join(self.root, "images", id)).convert('L')
            mask = Image.open(os.path.join(self.root, "masks", self.cls_name, id)).convert('L')
            img, mask = np.array(img)/255.0, np.array(mask)/255.0
        else: # train_u
            img = Image.open(os.path.join(self.u_root, id)).convert('L')
            img, mask = np.array(img)/255.0, np.zeros((img.size[1], img.size[0]), dtype=np.uint8)

        # 验证时直接返回原图和mask
        if self.mode == 'val':
            return torch.from_numpy(img).unsqueeze(0).float(), torch.from_numpy(mask).long()

        random_value = random.random()
        if random_value > 0.5:
            img, mask = random_flip(img, mask)
        if random.random() > 0.5:
            img, mask = random_rotate(img, mask)

        x, y = img.shape
        img = zoom(img, (self.size / x, self.size / y), order=0)
        mask = zoom(mask, (self.size / x, self.size / y), order=0)

        if self.mode == 'train_l':
            return torch.from_numpy(img).unsqueeze(0).float(), torch.from_numpy(np.array(mask)).long()

        img = Image.fromarray((img * 255).astype(np.uint8))
        img_s1, img_s2 = deepcopy(img), deepcopy(img)
        img = torch.from_numpy(np.array(img)).unsqueeze(0).float() / 255.0

        if random.random() < 0.8:
            img_s1 = transforms.ColorJitter(0.5, 0.5, 0.5, 0.25)(img_s1)
        img_s1 = blur(img_s1, p=0.5)
        cutmix_box1 = obtain_cutmix_box(self.size, p=0.5)
        img_s1 = torch.from_numpy(np.array(img_s1)).unsqueeze(0).float() / 255.0
        img_s1 = transforms.RandomErasing(p=0.5, scale=(0.02, 0.1))(img_s1)

        if random.random() < 0.8:
            img_s2 = transforms.ColorJitter(0.5, 0.5, 0.5, 0.25)(img_s2)
        img_s2 = blur(img_s2, p=0.5)
        cutmix_box2 = obtain_cutmix_box(self.size, p=0.5)
        img_s2 = torch.from_numpy(np.array(img_s2)).unsqueeze(0).float() / 255.0
        img_s2 = transforms.RandomErasing(p=0.5, scale=(0.02, 0.1))(img_s2)

        return img, img_s1, img_s2, cutmix_box1, cutmix_box2

    def __len__(self):
        return len(self.ids)

if __name__ == '__main__':
    dataset = CarotidArteryDataset('carotidartery', 'train_l', 'Artery', 
                                   'E:/qiuyi.ye/data/AIMP_20250424', 
                                   'T:/算法组/for秋意/现有可用视频', 256, 
                                   'splits/carotidartery/train.txt')
    img, mask = dataset[0]
    print(img.shape, mask.shape)

    dataset_v = CarotidArteryDataset('carotidartery', 'val', 'Artery', 
                                   'E:/qiuyi.ye/data/AIMP_20250424', 
                                   'T:/算法组/for秋意/现有可用视频', 256, 
                                   'splits/carotidartery/valid.txt')
    img, mask = dataset_v[0]
    print(img.shape, mask.shape)

    dataset_u = CarotidArteryDataset('carotidartery', 'train_u', 'Artery', 
                                   'E:/qiuyi.ye/data/AIMP_20250424', 
                                   'T:/算法组/for秋意/现有可用视频', 256, 
                                   'splits/carotidartery/unlabeled.txt')
    img, img_s1, img_s2, cutmix_box1, cutmix_box2 = dataset_u[0]
    print(img.shape, img_s1.shape, img_s2.shape, cutmix_box1.shape, cutmix_box2.shape)